import numpy as np
from mindspore import Tensor, export, load_checkpoint, load_param_into_net

import os, sys, argparse
sys.path.append(os.path.abspath(os.path.join(os.getcwd())))
path = os.path.abspath(os.path.join(os.getcwd()))
from mindspore.common import set_seed
import numpy as np
import mindspore 
from mindspore import load_checkpoint, load_param_into_net
from mind3d.models.PointTransformer import PointTransformerSeg
from mind3d.models.pointnet import PointNet_seg
from mind3d.tools.segmentation.shapenetpart.eval import ModifyNetwork

def Export(args_opt):
    path = os.path.abspath(os.path.join(os.getcwd(), "configs/", str(args.model_name)))
    print("input: ", path )
    if args_opt.model == "pointtransformer": 
        network = PointTransformerSeg()
        network = ModifyNetwork(network)
        file_name = "PointTransformerSeg"
        pretrain_ckpt_path = args_opt.ExportSeg.export_file
        file_format = "MINDIR" 
        input_shape = [args_opt.ExportSeg.batch_size, 1024, 22]
        load_param_into_net(network, load_checkpoint(pretrain_ckpt_path))
        input_array = Tensor(np.random.uniform(-1.0, 1.0, size=input_shape).astype(np.float32))
        export(network, input_array, file_name = path +"/"+ file_name, file_format = file_format)
    elif args_opt.model == "PointNet_seg":
        network = PointNet_seg()
        file_name = "PointNetSeg"
        pretrain_ckpt_path = args_opt.ExportSeg.export_file
        file_format = "MINDIR" 
        load_param_into_net(network, load_checkpoint(pretrain_ckpt_path))
        input_data = Tensor(np.random.uniform(0.0, 1.0, size=[32, 1024, 6]), mindspore.float32)
        input_label = Tensor(np.random.uniform(0.0, 1.0, size=[32, 1, 16]), mindspore.float32)
        export(network, input_data, input_label, file_name=path +"/"+ file_name, file_format=file_format)


    print(f"Successful export " + str(args_opt.model))


def main(args):
    from mind3d.utils.PointTransformerUtils import AttrDict,create_attr_dict
    import yaml
    #path = os.path.abspath(os.path.join(os.getcwd(), "configs/", str(args.model_name)))
    with open(args.opt, 'r') as f:
        args_ops = AttrDict(yaml.safe_load(f.read()))
    create_attr_dict(args_ops)
    Export(args_opt=args_ops)



if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Shapenet segmentation train.')
    parser.add_argument('-opt',default="./configs/pointtransformer/pointtransformerseg.yaml",
                        help='pointtransformer or pointnet')
    args = parser.parse_known_args()[0]
    main(args)
